import sys
import os
import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import random
import hashlib
import argparse
from tqdm import tqdm

label_colours = [
    [0, 0, 0],          # 黑色
    [255, 0, 0],        # 红色
    [255, 255, 0],      # 黄色
    [0, 255, 0],        # 绿色
    [128, 0, 128],      # 紫色
    [0, 255, 255],      # 青色
    [192, 192, 192],    # 银色
    [255, 0, 255],      # 品红
    [165, 42, 42],      # 棕色
    [128, 128, 128],    # 灰色
    [128, 0, 0],        # 栗色
    [250, 250, 210],    # 浅黄绿色
    [128, 128, 0],      # 橄榄色
    [255, 250, 240],    # 翡翠白
    [0, 128, 0],        # 深绿色
    [0, 128, 128],      # 水鸭色
    [0, 0, 128],        # 海军蓝
    [255, 165, 0],      # 橙色
    [255, 215, 0],      # 金色
    [173, 216, 230],    # 淡蓝色
    [240, 128, 128],    # 淡珊瑚色
    [219, 112, 147],    # 紫红色
    [255, 20, 147],     # 深粉色
    [255, 105, 180],    # 热粉色
    [255, 192, 203],    # 粉色
    [0, 0, 255],        # 蓝色
    [255, 245, 238],    # 海贝色
    [240, 255, 240],    # 蜜色
    [240, 248, 255],    # 爱丽丝蓝
    [248, 248, 255],    # 幽灵白
    [240, 255, 255],    # 天蓝色
    [255, 245, 238],    # 海贝色
    [245, 255, 250],    # 薄荷奶油
    [240, 255, 240],    # 蜜露色
    [190, 190, 190],    # 银灰色
    [135, 206, 235],    # 天蓝色
    [0, 191, 255],      # 深天蓝
    [30, 144, 255],     # 道奇蓝
    [250, 235, 215],    # 古代白色
    [245, 245, 220],    # 米色
    [255, 239, 213],    # 木瓜色
    [105, 105, 105],    # 弱灰色
    [255, 228, 196],    # 斑驳色
    [169, 169, 169],    # 暗灰色
    [211, 211, 211],    # 亮灰色
    [255, 218, 185],    # 桃色
    [255, 222, 173],    # 纳瓦霍白
    [47, 79, 79],       # 深石板灰
    [112, 128, 144],    # 石板灰
    [119, 136, 153],    # 浅石板灰
    [70, 130, 180],     # 钢蓝色
    [176, 196, 222],    # 浅钢蓝
    [100, 149, 237],    # 矢车菊蓝
    [65, 105, 225],     # 皇家蓝
    [0, 0, 205],        # 中蓝色
    [255, 228, 181],    # 种色
    [148, 0, 211],      # 深紫罗兰色
    [153, 50, 204],     # 深兰花色
    [186, 85, 211],     # 中兰花色
    [128, 0, 128],      # 紫色
    [216, 191, 216],    # 蓟色
    [221, 160, 221],    # 梅红色
    [255, 182, 193],    # 浅粉红色
    [0, 0, 139],        # 暗蓝色
    [255, 248, 220],    # 玉色
    [255, 250, 205],    # 柠檬薄纱
    [25, 25, 112],      # 午夜蓝
    [255, 250, 250],    # 雪色
    [72, 61, 139],      # 暗板岩蓝
    [106, 90, 205],     # 板岩蓝
    [123, 104, 238],    # 中紫罗兰色
    [147, 112, 219],    # 中紫色
    [138, 43, 226],     # 蓝紫色
    [75, 0, 130],       # 靛青
    [139, 0, 139],      # 暗紫色
    [238, 130, 238],    # 紫罗兰色
    [250, 128, 114],    # 鲑鱼色
    [255, 99, 71],      # 番茄色
    [255, 69, 0],       # 橙红色
    [255, 165, 0],      # 橙色
    [255, 245, 238],    # 海贝色
    [245, 255, 250],    # 薄荷奶油色
    [255, 0, 255],      # 洋红
    [255, 0, 255],      # 品红
    [218, 112, 214],    # 兰花色
    [199, 21, 133],     # 中紫红色
    [208, 32, 144],     # 紫红色
    [255, 20, 147],     # 深粉色
    [255, 105, 180],    # 热粉色
    [255, 182, 193],    # 浅粉红色
    [255, 192, 203],    # 粉红
    [219, 112, 147],    # 紫红色
    [255, 160, 122],    # 浅鲑鱼色
    [255, 140, 0],      # 暗橙色
    [233, 150, 122],    # 深鲑鱼色
    [240, 128, 128],    # 淡珊瑚色
    [205, 92, 92],      # 印第安红
    [220, 20, 60],      # 猩红
    [255, 0, 0],        # 红色
    [178, 34, 34],      # 火砖色
    [139, 0, 0],        # 暗红色
    [255, 250, 205],    # 柠檬薄纱色
    [255, 215, 0],      # 金色
    [184, 134, 11],     # 暗金色
    [218, 165, 32],     # 金麒麟色
    [238, 232, 170],    # 浅卡其色
    [189, 183, 107],    # 暗卡其色
    [255, 228, 181],    # 种色
    [253, 245, 230],    # 老花白
    [250, 240, 230],    # 亚麻色
    [0, 206, 209],      # 深绿宝石
    [255, 245, 238],    # 海贝色
    [240, 255, 240],    # 蜜露色
    [245, 255, 250],    # 薄荷奶油色
    [127, 255, 212],    # 绿玉色
    [255, 255, 240],    # 象牙色
    [0, 255, 127],      # 弱绿色
    [255, 218, 185],    # 桃色
    [255, 239, 213],    # 木瓜色
    [60, 179, 113],     # 中春绿色
    [240, 255, 255],    # 天蓝色
    [240, 248, 255],    # 爱丽丝蓝
    [0, 255, 255],      # 青色
    [255, 239, 213],    # 木瓜色
    [255, 218, 185],    # 桃色
    [240, 230, 140],    # 卡其色
    [245, 245, 245],    # 白烟
    [255, 250, 250],    # 雪色
    [255, 255, 0],      # 黄色
    [128, 128, 0],      # 橄榄色
    [255, 255, 224],    # 浅黄色
    [250, 250, 210],    # 浅卡其色
    [224, 255, 255],    # 淡青色
    [175, 238, 238],    # 绿宝石
    [0, 255, 0],        # 绿色
    [255, 228, 181],    # 种色
    [250, 240, 230],    # 亚麻色
    [64, 224, 208],     # 绿松石色
    [255, 222, 173],    # 纳瓦霍白
    [240, 230, 140],    # 卡其色
    [255, 255, 0],      # 黄色
    [255, 215, 0],      # 金色
    [238, 232, 170],    # 淡卡其色
    [240, 255, 240],    # 蜜露色
    [240, 255, 255],    # 天蓝色
    [240, 248, 255],    # 爱丽丝蓝
    [248, 248, 255],    # 幽灵白
    [245, 245, 245],    # 白烟
    [255, 250, 240],    # 翡翠白
    [72, 209, 204],     # 中绿宝石
    [250, 250, 210],    # 浅卡其色
    [32, 178, 170],     # 暗绿宝石
    [127, 255, 0],      # 弱绿色
    [124, 252, 0],      # 草绿色
    [0, 255, 0],        # 亮绿色
    [50, 205, 50],      # 黄绿色
    [152, 251, 152],    # 苍白绿色
    [143, 188, 143],    # 暗海绿色
    [0, 100, 0],        # 暗绿色
    [255, 222, 173],    # 纳瓦霍白
    [255, 250, 240],    # 翡翠白
    [253, 245, 230],    # 老花白
    [248, 248, 255],    # 幽灵白
    [0, 139, 139],      # 暗青色
    [0, 128, 128],      # 水鸭色
    [255, 255, 240],    # 象牙色
    [47, 79, 79],       # 深石板灰
    [0, 128, 128],      # 水鸭色
    [0, 250, 154],      # 春绿色
    [34, 139, 34],      # 森林绿
    [0, 128, 0],        # 绿色
    [107, 142, 35],     # 橄榄军服色
    [85, 107, 47],      # 暗橄榄绿色
    [128, 128, 0],      # 橄榄色
    [189, 183, 107],    # 暗卡其色
    [255, 250, 250],    # 雪色
    [255, 255, 255],    # 白色
]

def get_color(value):
    color = label_colours[value]
    return color

def display_mask_image(mask, void_values, show_image_width=720):
    """显示彩色图片并在左上角显示颜色对应的数值"""
    unique_values = np.unique(mask)
    colormap = {value: [0,0,0] if value in void_values else get_color(value) for value in unique_values}
    colored_image = np.zeros((mask.shape[0], mask.shape[1], 3), dtype=np.uint8)
    for value, color in colormap.items():
        colored_image[mask == value] = color
    scale = show_image_width / colored_image.shape[1]
    colored_image = cv2.resize(colored_image,(int(colored_image.shape[1]*scale), int(colored_image.shape[0]*scale)))
    font = cv2.FONT_HERSHEY_SIMPLEX
    font_scale = 0.5
    text_size, baseline = cv2.getTextSize('Value 255: [255,255,255]', fontFace=font, fontScale=font_scale, thickness=1)
    text_width = text_size[0]
    text_height = text_size[1] + 10
    y_offset = 20
    colored_image[:text_height*len(colormap)+y_offset, :text_width] = [255,255,255]
    for value, color in colormap.items():
        cv2.putText(colored_image,  f'Value {value}: {color}', (10, y_offset), fontFace=font, fontScale=font_scale, color=color, thickness=2)
        y_offset += text_height
    cv2.imshow('Mask Images', colored_image)
    

def main():
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description="mask show")
    parser.add_argument("filepath", help="文件的路径")
    parser.add_argument("-v", "--void_values", nargs='+', type=int, default=[], help="void mask values, if set this list mask value will show black")
    parser.add_argument("-s", "--suffixes", nargs='+', default=['jpg', 'jpeg', 'png'], type=str, help="image file suffixes")
    parser.add_argument("-w", "--show_width", type=int, default=720, help="show image width")
    args = parser.parse_args()
    image_paths = []
    if os.path.isdir(args.filepath):
        image_paths = [os.path.join(args.filepath, f) for f in os.listdir(args.filepath) if f.endswith(tuple(args.suffixes))]
    elif args.filepath.endswith(tuple(args.suffixes)):
        image_paths.append(args.filepath)
    cv2.namedWindow('Mask Images')
    for path in tqdm(image_paths):
        # 读取单通道8位的掩码图片
        mask = Image.open(path).convert('L')
        mask = np.array(mask)
        cv2.setWindowTitle('Mask Images', path)
        # 显示彩色图片
        display_mask_image(mask, args.void_values, args.show_width)
        key = cv2.waitKey(0)
        if key == 27:  # 按下esc,退出
            break
    cv2.destroyAllWindows()

# 使用示例
if __name__ == "__main__":
    main()